package org.zoomdev.zoom.retry;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Consumer;

/**
 * 增强版重试服务 - 支持Job执行
 */
public class EnhancedRetryService {

    // 核心Timer实例
    private final DriftProofTimer timer;

    // 任务ID生成器
    private final AtomicLong taskIdGenerator = new AtomicLong(0);

    // 任务注册表
    private final Map<Long, RetryJob> registeredJobs = new ConcurrentHashMap<>();

    public EnhancedRetryService() {
        this.timer = new DriftProofTimer();
    }

    public EnhancedRetryService(DriftProofTimer timer) {
        this.timer = timer;
    }

    /**
     * 使用示例
     */
    public static void main(String[] args) throws InterruptedException {
        // 创建重试服务
        EnhancedRetryService retryService = new EnhancedRetryService();
        retryService.start();

        // 创建调用器实现
        EnhancedRetryService.Invoker invoker = (target, data, retryCount) -> {
            System.out.printf("Executing job on target: %s with data: %s (retry #%d)%n",
                    target, data, retryCount);

            // 模拟失败
            if (retryCount < 2 && Math.random() > 0.3) {
                throw new RuntimeException("Job execution failed");
            }
        };

        // 创建Job实例
        EnhancedRetryService.Job job = new EnhancedRetryService.Job(
                invoker,
                "TargetService",
                "PayloadData"
        );

        // 安排带重试的Job任务（简单重试策略）
        long jobTaskId = retryService.scheduleJob(
                job,
                1, TimeUnit.SECONDS,
                new SimpleRetryStrategy(500, TimeUnit.MILLISECONDS),
                10L, TimeUnit.SECONDS, // 最大重试时间
                5 // 最大重试次数
        );

        // 监控任务状态
        ScheduledExecutorService monitor = Executors.newScheduledThreadPool(1);
        monitor.scheduleAtFixedRate(() -> {
            System.out.println("\n=== Job Status ===");
            System.out.println("Job execution count: " + job.getRetryCount());
            System.out.println("Pending jobs: " + retryService.registeredJobs.size());
            System.out.println("Timer pending tasks: " + retryService.getTimer().pendingTasks());

            retryService.getTimer().getNextExpirationTime().ifPresent(next -> {
                System.out.println("Next job scheduled in: " + (next - System.currentTimeMillis()) + "ms");
            });
        }, 0, 2, TimeUnit.SECONDS);

        // 运行20秒
        Thread.sleep(20000);

        // 停止服务
        System.out.println("\nStopping retry service...");
        retryService.stop();
        monitor.shutdown();

        System.out.println("Service stopped. Program exiting.");
    }

    /**
     * 启动服务
     */
    public void start() {
        timer.start();
    }

    /**
     * 停止服务
     */
    public void stop() {
        timer.stop();
    }

    /**
     * 获取底层Timer实例
     */
    public DriftProofTimer getTimer() {
        return timer;
    }

    /**
     * 安排Job任务（一次性）
     */
    public long scheduleJob(Job job, long delay, TimeUnit unit,
                            RetryStrategy retryStrategy, Long maxRetryDuration, TimeUnit maxUnit,
                            Integer maxRetries) {
        RetryJob retryJob = new RetryJob(job, retryStrategy,
                maxRetryDuration != null ? maxUnit.toMillis(maxRetryDuration) : -1,
                maxRetries);
        return scheduleJobTask(retryJob, delay, unit, false, -1, null);
    }

    /**
     * 安排周期性Job任务
     */
    public long scheduleJobWithFixedDelay(Job job, long initialDelay, long period, TimeUnit unit,
                                          RetryStrategy retryStrategy, Long maxRetryDuration, TimeUnit maxUnit,
                                          Integer maxRetries) {
        RetryJob retryJob = new RetryJob(job, retryStrategy,
                maxRetryDuration != null ? maxUnit.toMillis(maxRetryDuration) : -1,
                maxRetries);
        return scheduleJobTask(retryJob, initialDelay, unit, true, period, unit);
    }

    /**
     * 取消任务
     */
    public void cancel(long taskId) {
        registeredJobs.remove(taskId);
        timer.cancel(taskId);
    }

    /**
     * 检查任务是否已取消
     */
    public boolean isCanceled(long taskId) {
        return timer.isCanceled(taskId);
    }

    /**
     * 设置全局异常处理器
     */
    public void setExceptionHandler(Consumer<Exception> handler) {
        timer.setExceptionHandler(handler);
    }

    /**
     * 安排Job任务到Timer
     */
    private long scheduleJobTask(RetryJob retryJob, long delay, TimeUnit unit,
                                 boolean periodic, long period, TimeUnit periodUnit) {
        long taskId = taskIdGenerator.incrementAndGet();
        registeredJobs.put(taskId, retryJob);

        // 封装Job执行逻辑
        Runnable jobRunner = () -> executeJobWithRetry(retryJob, taskId);

        // 委托给Timer安排任务
        if (periodic) {
            return timer.scheduleWithFixedDelay(jobRunner, delay, period, unit);
        } else {
            return timer.schedule(jobRunner, delay, unit);
        }
    }

    /**
     * 执行Job并处理重试逻辑
     */
    private void executeJobWithRetry(RetryJob retryJob, long taskId) {
        // 记录初始执行时间
        if (retryJob.firstFailureTime == 0) {
            retryJob.firstFailureTime = System.currentTimeMillis();
        }

        try {
            // 设置当前重试计数并执行Job
            retryJob.job.retryCount = retryJob.currentRetryCount;
            retryJob.job.run();

            // 执行成功，重置重试状态
            retryJob.resetRetryState();
        } catch (Exception e) {
            handleJobFailure(retryJob, taskId, e);
        }
    }

    /**
     * 处理Job失败情况
     */
    private void handleJobFailure(RetryJob retryJob, long taskId, Exception e) {
        retryJob.exceptionHandler.accept(e);

        // 检查是否超过最大重试次数
        if (retryJob.exceedsMaxRetries()) {
            System.out.printf("Job %d exceeded max retries (%d). Stopping retries.%n",
                    taskId, retryJob.maxRetries);
            retryJob.retryAborted = true;
            return;
        }

        // 检查是否超过最大重试时间
        if (retryJob.exceedsMaxRetryDuration()) {
            System.out.printf("Job %d exceeded max retry duration (%dms). Stopping retries.%n",
                    taskId, retryJob.maxRetryDuration);
            retryJob.retryAborted = true;
            return;
        }

        // 计算下次重试延迟
        long nextDelay = retryJob.retryStrategy.calculateNextDelay(retryJob.currentRetryCount);
        if (nextDelay < 0) {
            System.out.printf("Job %d has no more retry attempts. Stopping retries.%n", taskId);
            retryJob.retryAborted = true;
            return;
        }

        // 更新重试计数
        retryJob.currentRetryCount++;
        System.out.printf("Scheduling retry #%d for job %d after %dms%n",
                retryJob.currentRetryCount, taskId, nextDelay);

        // 使用Timer调度重试
        timer.schedule(() -> executeJobWithRetry(retryJob, taskId), nextDelay, TimeUnit.MILLISECONDS);
    }

    /**
     * 调用器接口
     */
    @FunctionalInterface
    public interface Invoker {
        void invoke(Object target, Object data, int retryCount);
    }

    /**
     * 重试策略接口
     */
    public interface RetryStrategy {
        /**
         * 计算下一次重试的延迟
         *
         * @param currentRetry 当前重试次数（从0开始）
         * @return 下一次重试的延迟时间（毫秒）
         */
        long calculateNextDelay(int currentRetry);
    }

    /**
     * Job重试包装类
     */
    private static class RetryJob {
        private final Job job;
        private final RetryStrategy retryStrategy;
        private final long maxRetryDuration; // 最大重试时间(ms)
        private final Integer maxRetries;    // 最大重试次数
        private long firstFailureTime;        // 第一次失败时间戳(ms)
        private int currentRetryCount;        // 当前重试次数
        private boolean retryAborted;        // 是否终止重试
        private Consumer<Exception> exceptionHandler = e -> {
        };

        public RetryJob(Job job, RetryStrategy retryStrategy,
                        long maxRetryDuration, Integer maxRetries) {
            this.job = job;
            this.retryStrategy = retryStrategy;
            this.maxRetryDuration = maxRetryDuration;
            this.maxRetries = maxRetries;
        }

        public void resetRetryState() {
            firstFailureTime = 0;
            currentRetryCount = 0;
            retryAborted = false;
        }

        public boolean exceedsMaxRetryDuration() {
            if (maxRetryDuration <= 0) return false;
            return firstFailureTime > 0 &&
                    (System.currentTimeMillis() - firstFailureTime) > maxRetryDuration;
        }

        public boolean exceedsMaxRetries() {
            if (maxRetries == null) return false;
            return currentRetryCount >= maxRetries;
        }

        public void setExceptionHandler(Consumer<Exception> handler) {
            this.exceptionHandler = handler;
        }
    }

    /**
     * 任务接口
     */
    public static class Job {
        private int retryCount;
        private Object data;
        private Object target;
        private Invoker invoker;

        public Job(Invoker invoker, Object target, Object data) {
            this.invoker = invoker;
            this.target = target;
            this.data = data;
        }

        public void run() {
            invoker.invoke(target, data, retryCount);
        }

        // 访问器方法
        public int getRetryCount() {
            return retryCount;
        }

        public Object getData() {
            return data;
        }

        public Object getTarget() {
            return target;
        }
    }

    /**
     * 简单重试策略（固定延迟）
     */
    public static class SimpleRetryStrategy implements RetryStrategy {
        private final long fixedDelay; // 固定延迟(ms)

        public SimpleRetryStrategy(long delay, TimeUnit unit) {
            this.fixedDelay = unit.toMillis(delay);
        }

        @Override
        public long calculateNextDelay(int currentRetry) {
            return fixedDelay;
        }
    }

    /**
     * 延时序列重试策略
     */
    public static class DelaySequenceRetryStrategy implements RetryStrategy {
        private final long[] delays; // 延迟序列(ms)

        public DelaySequenceRetryStrategy(long[] delays, TimeUnit unit) {
            this.delays = Arrays.copyOf(delays, delays.length);
            for (int i = 0; i < this.delays.length; i++) {
                this.delays[i] = unit.toMillis(delays[i]);
            }
        }

        public DelaySequenceRetryStrategy(List<Long> delays, TimeUnit unit) {
            this.delays = delays.stream().mapToLong(unit::toMillis).toArray();
        }

        @Override
        public long calculateNextDelay(int currentRetry) {
            if (currentRetry >= delays.length) {
                return -1; // 超出序列长度
            }
            return delays[currentRetry];
        }
    }

    /**
     * 定时重试策略（固定间隔）
     */
    public static class ScheduledRetryStrategy implements RetryStrategy {
        private final long fixedInterval; // 固定间隔(ms)

        public ScheduledRetryStrategy(long interval, TimeUnit unit) {
            this.fixedInterval = unit.toMillis(interval);
        }

        @Override
        public long calculateNextDelay(int currentRetry) {
            return fixedInterval;
        }
    }
}